import os
import random
import shutil

# 配置路径
dataset_path = r"E:\rune\tl official"  # 替换为你的数据集路径
output_path = r"E:\rune\tl official\a" # 设定输出文件夹路径
train_ratio = 0.75  # 训练集比例
val_ratio = 0.25 # 验证集比例，测试集比例自动计算为 1 - train_ratio - val_ratio

# 数据集划分文件夹
splits = ["train", "valid", "test"]
for split in splits:
    os.makedirs(os.path.join(output_path, split, "images"), exist_ok=True)
    os.makedirs(os.path.join(output_path, split, "labels"), exist_ok=True)

# 获取所有图片文件（假设图片格式为 png 或 jpg）
image_files = sorted([f for f in os.listdir(dataset_path) if f.endswith((".png", ".jpg"))])
random.shuffle(image_files)  # 随机打乱数据

# 计算划分数量
total_count = len(image_files)
train_count = int(total_count * train_ratio)
val_count = int(total_count * val_ratio)
test_count = total_count - train_count - val_count

# 数据分配
data_splits = {
    "train": image_files[:train_count],
    "valid": image_files[train_count:train_count + val_count],
    "test": image_files[train_count + val_count:]
}

# 复制文件到对应文件夹
for split, images in data_splits.items():
    for img in images:
        base_name = os.path.splitext(img)[0]  # 获取文件名（无扩展名）
        img_src = os.path.join(dataset_path, img)
        label_src = os.path.join(dataset_path, base_name + ".txt")

        img_dst = os.path.join(output_path, split, "images", img)
        label_dst = os.path.join(output_path, split, "labels", base_name + ".txt")

        shutil.copy(img_src, img_dst)  # 复制图片
        if os.path.exists(label_src):  # 确保标注文件存在再复制
            shutil.copy(label_src, label_dst)

    print(f"{split} 集合: {len(images)} 张图片")

print("数据集划分完成！")